package com.hky.aiagent.rag;

import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;

/**
 * 创建自定义 rag 检索增强工厂
 *
 * @author hky
 */

public class LoveAppRagCustomAdvisorFactory {

    /**
     * 创建一个Advisor实例，该实例用于根据给定的状态筛选向量存储中的文档，并按相关性排序
     *
     * @param vectorStore 向量存储，用于文档的检索
     * @param status 状态字符串，用于筛选符合条件的文档
     * @return 返回一个初始化的Advisor对象，该对象使用了特定配置的文档检索器
     */
    public static Advisor createAdvisor(VectorStore vectorStore, String status) {
        // 创建一个过滤表达式，用于匹配状态字段与给定状态值相同的文档
        Filter.Expression expressionBuilder = new FilterExpressionBuilder()
                .eq("status", status)
                .build();

        // 构建一个文档检索器，配置向量存储、过滤条件、相似度阈值和返回的文档数量
        DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .filterExpression(expressionBuilder) //过滤条件
                .similarityThreshold(0.5) // 相似度阈值
                .topK(3) // 返回的文档数量
                .build();

        // 构建一个检索增强型顾问，配置文档检索器
        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(documentRetriever)
                .queryAugmenter(AppContextualQueryAugmenterFactory.createInstance())
                .build();
    }

}
